import numpy as np


datas = [
    {"city": "吕梁", "area": 80, "rooms": 2, "school": 1, "style": 0, "price": 5000},
    {"city": "太原", "area": 100, "rooms": 3, "school": 0, "style": 1, "price": 7000},
    {"city": "吕梁", "area": 120, "rooms": 4, "school": 1, "style": 1, "price": 9000},
    {"city": "太原", "area": 90, "rooms": 2, "school": 0, "style": 0, "price": 6000},
    {"city": "吕梁", "area": 75, "rooms": 3, "school": 1, "style": 0, "price": 5500},
    {"city": "太原", "area": 110, "rooms": 3, "school": 1, "style": 1, "price": 8000},
    {"city": "吕梁", "area": 95, "rooms": 2, "school": 0, "style": 1, "price": 6500},
    {"city": "太原", "area": 85, "rooms": 4, "school": 0, "style": 0, "price": 7500},
    {"city": "吕梁", "area": 65, "rooms": 1, "school": 1, "style": 0, "price": 4500},
    {"city": "太原", "area": 105, "rooms": 3, "school": 1, "style": 1, "price": 8500},
    {"city": "晋中", "area": 130, "rooms": 4, "school": 1, "style": 1, "price": 9500},
    {"city": "晋中", "area": 75, "rooms": 2, "school": 0, "style": 0, "price": 5500},
    {"city": "晋中", "area": 95, "rooms": 3, "school": 1, "style": 0, "price": 6500},
    {"city": "晋中", "area": 110, "rooms": 4, "school": 0, "style": 1, "price": 8000},
    {"city": "晋中", "area": 85, "rooms": 2, "school": 1, "style": 1, "price": 7000}
]

X = []  # 特征值
Y = []  # 标签值
city_mark = {"吕梁": 1, "太原": 2, "晋中": 3}  # 城市标记

for item in datas:
    single = []
    single.append(city_mark[item["city"]])
    single.append(item["area"])
    single.append(item["rooms"])
    single.append(item["school"])
    single.append(item["style"])
    X.append(single)
    Y.append(item["price"])

X = np.array(X)
Y = np.array(Y)

# 添加偏置项
X = np.hstack((np.ones((X.shape[0], 1)), X))

# 计算theta
theta = np.linalg.pinv(X.T.dot(X)).dot(X.T).dot(Y)

# 预测示例
input_features = np.array([3, 2, 150, 2, 1, 1])  # 假设第一个元素是偏置项
prediction = theta.dot(input_features)
print(prediction)
